"""Vulnerability merge model

Revision ID: 526aa91cac98
Revises: 2db31733fb78
Create Date: 2019-07-01 19:19:10.166057+00:00

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = '526aa91cac98'
down_revision = '085188e0a016'
branch_labels = None
depends_on = None


def upgrade():
    conn = op.get_bind()
    op.create_table(
        'knowledge_base',
        sa.Column('id', sa.Integer, primary_key=True),
        sa.Column('vulnerability_template_id', sa.Integer),
        sa.Column('faraday_kb_id', sa.Text, nullable=False),
        sa.Column('reference_id', sa.Integer, nullable=False),
        sa.Column('external_identifier', sa.Text),
        sa.Column('script_name', sa.Text),
        sa.Column('tool_name', sa.Text),
        sa.Column('false_positive', sa.Integer, nullable=False, default=0),
        sa.Column('verified', sa.Integer, nullable=False, default=0),
        sa.UniqueConstraint('external_identifier', 'tool_name', 'reference_id',
                            name='uix_externalidentifier_toolname_referenceid')
    )

    op.create_foreign_key(
        'knowledge_base_vulnerability_template_id_fkey', 'knowledge_base',
        'vulnerability_template', ['vulnerability_template_id'], ['id']
    )

    op.add_column('vulnerability',
                  sa.Column(
                      'association_date',
                      sa.DateTime(), nullable=True
                  )
                  )

    op.add_column('vulnerability',
                  sa.Column(
                      'vulnerability_template_id',
                      sa.Integer(),
                      nullable=True
                  )
                  )

    op.create_foreign_key(
        'vulnerability_vulnerability_template_id_fkey',
        'vulnerability',
        'vulnerability_template', ['vulnerability_template_id'], ['id'],
        ondelete="RESTRICT"
    )

    op.add_column('vulnerability',
                  sa.Column(
                      'vulnerability_duplicate_id',
                      sa.Integer(),
                      nullable=True
                  )
                  )

    op.add_column('vulnerability',
                  sa.Column(
                      'disassociated_manually',
                      sa.Boolean(),
                      nullable=False,
                      server_default='false',
                  )
                  )

    op.create_foreign_key(
        'vulnerability_vulnerability_duplicate_id_fkey',
        'vulnerability',
        'vulnerability', ['vulnerability_duplicate_id'], ['id']
    )

    op.add_column('vulnerability_template',
                  sa.Column(
                      'shipped',
                      sa.Boolean(),
                      nullable=False,
                      server_default='false',
                  )
                  )

    conn.execute('ALTER TABLE vulnerability_template DROP CONSTRAINT uix_vulnerability_template_name')
    conn.execute(
        'ALTER TABLE vulnerability_template ADD CONSTRAINT uix_vulnerability_template_name UNIQUE (name, shipped)')


def downgrade():
    conn = op.get_bind()

    conn.execute('ALTER TABLE vulnerability_template DROP CONSTRAINT uix_vulnerability_template_name')
    op.drop_column('vulnerability_template', 'shipped')
    conn.execute('ALTER TABLE vulnerability_template ADD CONSTRAINT uix_vulnerability_template_name UNIQUE (name)')

    op.drop_table('knowledge_base')

    op.drop_column('vulnerability', 'vulnerability_duplicate_id')
    op.drop_column('vulnerability', 'vulnerability_template_id')
    op.drop_column('vulnerability', 'association_date')
    op.drop_column('vulnerability', 'disassociated_manually')
